package com.alibaba.datax.plugin.reader.mongodbreader;

import java.util.List;

import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.plugin.RecordSender;
import com.alibaba.datax.common.spi.Reader;
import com.alibaba.datax.common.util.Configuration;
import com.alibaba.datax.plugin.reader.mongodbreader.util.CollectionSplitUtil;
import com.alibaba.datax.plugin.reader.mongodbreader.util.MongoClientUtil;
import com.alibaba.datax.plugin.reader.mongodbreader.util.MongoUtil;
import com.google.common.base.Strings;
import com.mongodb.MongoClient;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import org.bson.Document;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Created by jianying.wcj on 2015/3/19 0019.
 * Modified by mingyan.zc on 2016/6/13.
 * Modified by mingyan.zc on 2017/7/5.
 */
public class MongoDBReader extends Reader {

    public static class Job extends Reader.Job {
        private static final Logger LOG = LoggerFactory.getLogger(Job.class);
        private Configuration originalConfig = null;

        private MongoClient mongoClient;

        @Override
        public List<Configuration> split(int adviceNumber) {
            return CollectionSplitUtil.doSplit(originalConfig, adviceNumber, mongoClient);
        }

        @Override
        public void init() {
            this.originalConfig = super.getPluginJobConf();
            String userName = originalConfig.getString(KeyConstant.MONGO_USER_NAME);
            String password = originalConfig.getString(KeyConstant.MONGO_USER_PASSWORD);
            originalConfig.getNecessaryValue(KeyConstant.MONGO_DB_NAME, MongoDBReaderErrorCode.DB_NAME_IS_NULL);
            originalConfig.getNecessaryValue(KeyConstant.MONGO_COLLECTION_NAME, MongoDBReaderErrorCode.COLLECTION_NAME_IS_NULL);
            originalConfig.getNecessaryValue(KeyConstant.MONGO_COLUMN, MongoDBReaderErrorCode.COLUMN_IS_NULL);
            originalConfig.getNecessaryValue(KeyConstant.MONGO_ADDRESS, MongoDBReaderErrorCode.ADDRESS_IS_NULL);
            String database = originalConfig.getString(KeyConstant.MONGO_DB_NAME);
            String authDb = originalConfig.getString(KeyConstant.MONGO_AUTHDB, database);
            if (!Strings.isNullOrEmpty(userName) && !Strings.isNullOrEmpty(password)) {
                LOG.debug("JOB创建Mongo客户端");
                this.mongoClient = MongoClientUtil.initCredentialMongoClient(originalConfig, userName, password, authDb);
            } else {
                LOG.debug("JOB创建Mongo客户端");
                this.mongoClient = MongoClientUtil.initMongoClient(originalConfig);
            }
        }

        @Override
        public void destroy() {
            if(null != this.mongoClient){
                LOG.debug("JOB关闭Mongo客户端");
                this.mongoClient.close();
            }
        }
    }


    public static class Task extends Reader.Task {
        private static final Logger LOG = LoggerFactory.getLogger(Task.class);
        private Configuration readerSliceConfig;

        private MongoClient mongoClient;

        private String database = null;
        private String collection = null;

        @Override
        public void startRead(RecordSender recordSender) {
            MongoDatabase db = mongoClient.getDatabase(database);
            MongoCollection col = db.getCollection(this.collection);
            Document filter = MongoUtil.parseQueryFilter(readerSliceConfig);
            MongoCursor<Document> dbCursor = col.find(filter).iterator();
            MongoUtil.sendToWrite(readerSliceConfig, recordSender, dbCursor);
        }

        @Override
        public void init() {
            this.readerSliceConfig = super.getPluginJobConf();
            String userName = readerSliceConfig.getString(KeyConstant.MONGO_USER_NAME);
            String password = readerSliceConfig.getString(KeyConstant.MONGO_USER_PASSWORD);
            this.database = readerSliceConfig.getString(KeyConstant.MONGO_DB_NAME);
            String authDb = readerSliceConfig.getString(KeyConstant.MONGO_AUTHDB, this.database);
            if (!Strings.isNullOrEmpty(userName) && !Strings.isNullOrEmpty(password)) {
                LOG.error("TASK创建Mongo客户端");
                mongoClient = MongoClientUtil.initCredentialMongoClient(readerSliceConfig, userName, password, authDb);
            } else {
                LOG.error("TASK创建Mongo客户端");
                mongoClient = MongoClientUtil.initMongoClient(readerSliceConfig);
            }
            this.collection = readerSliceConfig.getString(KeyConstant.MONGO_COLLECTION_NAME);
            if (mongoClient == null || database == null || collection == null) {
                throw DataXException.asDataXException(MongoDBReaderErrorCode.ILLEGAL_VALUE,
                        MongoDBReaderErrorCode.ILLEGAL_VALUE.getDescription());
            }
        }

        @Override
        public void destroy() {
            if(null != this.mongoClient){
                LOG.error("TASK关闭Mongo客户端");
                this.mongoClient.close();
            }
        }

    }
}
